# genetic_algorithm.py
import random
import numpy as np
from typing import List
from vrp_model import VRPTWInstance, Customer


class GeneticAlgorithm:
    def __init__(self, instance: VRPTWInstance, pop_size=100, elite_rate=0.1,
                 crossover_rate=0.8, mutation_rate=0.1, max_generations=200):
        self.instance = instance
        self.pop_size = pop_size
        self.elite_rate = elite_rate
        self.crossover_rate = crossover_rate
        self.mutation_rate = mutation_rate
        self.max_generations = max_generations

    def initialize_population(self) -> List[List[int]]:
        population = []
        customers = [c.id for c in self.instance.customers if c.id != 0]
        for _ in range(self.pop_size):
            chromo = customers.copy()
            random.shuffle(chromo)
            population.append(chromo)  # 现在存储的是未分割的客户排列
        return population

    def split_routes(self, chromo: List[int]) -> List[List[int]]:
        routes = []
        current_route = [0]
        current_load = 0
        current_volume = 0
        current_dist = 0

        for gene in chromo:
            customer = next(c for c in self.instance.customers if c.id == gene)

            new_dist = current_dist + self.instance.distance_matrix[current_route[-1]][gene]
            return_dist = self.instance.distance_matrix[gene][0]

            if (current_load + customer.demand > self.instance.vehicle.capacity or
                    current_volume + customer.volume > self.instance.vehicle.volume or
                    (new_dist + return_dist) > self.instance.vehicle.max_range):

                current_route.append(0)
                routes.append(current_route)
                current_route = [0, gene]
                current_load = customer.demand
                current_volume = customer.volume
                current_dist = self.instance.distance_matrix[0][gene]
            else:
                current_route.append(gene)
                current_load += customer.demand
                current_volume += customer.volume
                current_dist = new_dist

        if len(current_route) > 1:
            current_route.append(0)
            routes.append(current_route)
        return routes

    def calculate_fitness(self, chromo: List[int]) -> float:
        routes = self.split_routes(chromo)
        total_cost = 0
        for route in routes:
            if len(route) < 3:
                continue

            total_cost += self.instance.fixed_cost
            distance = sum(self.instance.distance_matrix[i][j] for i, j in zip(route, route[1:]))

            if distance > self.instance.vehicle.max_range:
                return float('inf')

            total_cost += distance * self.instance.transport_cost

            current_time = 0
            for i in range(1, len(route) - 1):
                from_node = route[i - 1]
                to_node = route[i]
                current_time += self.instance.time_matrix[from_node][to_node]

                customer = next(c for c in self.instance.customers if c.id == to_node)
                if current_time < customer.early:
                    penalty = (customer.early - current_time) * self.instance.early_penalty
                elif current_time > customer.late:
                    penalty = (current_time - customer.late) * self.instance.late_penalty
                else:
                    penalty = 0
                total_cost += penalty

                current_time += 0  # 假设服务时间为0

            charge_time = (
                                      self.instance.vehicle.max_range - distance) / self.instance.vehicle.max_range * self.instance.charge_time
            total_cost += charge_time * (self.instance.charge_cost_rate + self.instance.opportunity_cost_rate)

        return total_cost

    def pmx_crossover(self, parent1, parent2):
        size = len(parent1)
        p1, p2 = sorted(random.sample(range(size), 2))

        # 创建子代副本
        child1 = parent1.copy()
        child2 = parent2.copy()

        # 建立映射关系
        mapping1 = {}
        mapping2 = {}
        for i in range(p1, p2 + 1):
            mapping1[parent2[i]] = parent1[i]
            mapping2[parent1[i]] = parent2[i]

        # 处理交叉区域外的冲突
        for i in list(range(0, p1)) + list(range(p2 + 1, size)):
            # 处理第一个子代
            while child1[i] in mapping1:
                child1[i] = mapping1[child1[i]]
            # 处理第二个子代
            while child2[i] in mapping2:
                child2[i] = mapping2[child2[i]]

        return child1, child2

    def mutate(self, chromo):
        if random.random() < self.mutation_rate:
            i, j = random.sample(range(len(chromo)), 2)
            chromo[i], chromo[j] = chromo[j], chromo[i]
        return chromo

    def local_search(self, chromo):
        best_chromo = chromo.copy()
        best_cost = self.calculate_fitness(chromo)

        for _ in range(5):
            new_chromo = chromo.copy()
            i, j = random.sample(range(len(new_chromo)), 2)
            new_chromo[i], new_chromo[j] = new_chromo[j], new_chromo[i]

            new_cost = self.calculate_fitness(new_chromo)
            if new_cost < best_cost:
                best_chromo = new_chromo
                best_cost = new_cost

        return best_chromo

    def solve(self):
        # 修改后的主循环逻辑
        population = self.initialize_population()
        best_solution = None
        best_fitness = float('inf')
        fitness_history = []

        for gen in range(self.max_generations):
            # 计算适应度时自动分割路径
            fitness = [self.calculate_fitness(chromo) for chromo in population]

            elite_size = int(self.elite_rate * self.pop_size)
            elite_indices = np.argsort(fitness)[:elite_size]
            elites = [population[i] for i in elite_indices]

            selected = []
            inv_fitness = [1 / f if f != 0 else 1e-6 for f in fitness]
            total = sum(inv_fitness)
            probabilities = [f / total for f in inv_fitness]

            for _ in range(self.pop_size - elite_size):
                selected.append(population[random.choices(range(len(population)), weights=probabilities)[0]])

            new_population = elites.copy()
            while len(new_population) < self.pop_size:
                parent1 = random.choice(selected)
                parent2 = random.choice(selected)

                if random.random() < self.crossover_rate:
                    child1, child2 = self.pmx_crossover(parent1, parent2)
                else:
                    child1, child2 = parent1.copy(), parent2.copy()

                child1 = self.mutate(child1)
                child2 = self.mutate(child2)
                new_population.extend([child1, child2])

            population = new_population[:self.pop_size]

            if gen % 10 == 0:
                population[0] = self.local_search(population[0])

            current_best = min(fitness)
            if current_best < best_fitness:
                best_fitness = current_best
                best_solution = population[np.argmin(fitness)]

            fitness_history.append(best_fitness)
            print(f"Generation {gen + 1}: Best Cost = {best_fitness:.2f}元")

        return best_solution, fitness_history